# given a set of redshifts, ras and decs, plots a pie plot.
import pylab as pl
import numpy as np
from cosmology import Cosmology,PC
from astro.constants import wlya,Ckms
from numpy import sin,cos
import q0107

RAD_PER_DEG = np.pi/180.
# minimum line detection significance.  Used for real and randoms.
NSIGMA = 3.

# pl.rc('text',usetex=True)
# pl.rc('axes',labelsize=16)
# pl.rc('xtick',labelsize=14)
# pl.rc('ytick',labelsize=14)

cosmo = Cosmology(H0=70,Om=0.3,Ol=0.7)
const = 1. / (1.e6 * PC) * cosmo.h

#wranges = np.rec.fromrecords([(1220,1523), (1621,2500)],names='mi,ma')
wranges = np.rec.fromrecords([(1220,1523), (0,2500)],names='mi,ma')
wlosmin = [cosmo.Dc(z)*const for z in (wranges.mi/wlya - 1)]
wlosmax = [cosmo.Dc(z)*const for z in (wranges.ma/wlya - 1)]

print wlosmin,wlosmax

def pie(gal, qsos, mode='ra', deg=False, zmin=None, zmax=None,
        magnify_theta=1., metals=False, legend=True, MAXSEP=25):
    """ Call with a galaxy record array (with ra, dec in degrees and z)
    and qso record array (ra, dec in degrees and z).  Makes a pie plot
    with galaxy LOS distance plotted as the r variable, and angular
    separation in degrees from the second QSO plotted as theta
    variable. Also plots the position of the QSO sight-lines.
    """

    #if lines is None:
    #    la = lb = lc = None    
    mode = mode.lower()
    if mode not in ('ra','dec'):
        raise Exception('Must specify ra or dec')
    #pl.ioff()
    fig = pl.figure()
    ax = fig.add_axes((0,0,1,1))

    if zmin is None:
        zmin = min(qsos.z.min(),gal.z.min())
    if zmax is None:
        zmax = max(qsos.z.max(),gal.z.max())
    
    # Use RA/dec of QSO 2 as reference angle.
    gal = gal[(zmin<=gal.z)&(gal.z<=zmax)]
    #qsos = qsos[(zmin<=qsos.z)&(qsos.z<=zmax)]
    gtheta = gal.ra if mode == 'ra' else gal.dec
    qtheta = qsos.ra if mode == 'ra' else qsos.dec
    #theta0 = qtheta[1]
    theta0 = qtheta[0]
    
    # exaggerate angular separations by the given factor
    sepdeg = (gtheta - theta0) * magnify_theta
    maxsep = min(MAXSEP, max(np.abs(sepdeg))*1.2)

    #remove galaxies outside maximum separation
    gal = gal[sepdeg < MAXSEP]
    sepdeg = sepdeg[sepdeg < MAXSEP]
    
    # convert redshifts to (distance * h) in units of Mpc
    losdist = np.array([cosmo.Dc(z) * const for z in gal.z])

    # plot axis limits:
    maxr = cosmo.Dc(zmax) * const
    minr = cosmo.Dc(zmin) * const
    x0 = cos(maxsep * RAD_PER_DEG) * minr
    y0 = sin(maxsep * RAD_PER_DEG) * minr
    x1 = cos(maxsep * RAD_PER_DEG) * maxr
    y1 = sin(maxsep * RAD_PER_DEG) * maxr
    ax.plot([x0,x1],[y0,y1],'k',lw=1.5,label=None)
    ax.plot([x0,x1],[-y0,-y1],'k',lw=1.5,label=None)
    r = np.linspace(maxsep * RAD_PER_DEG,-maxsep * RAD_PER_DEG,100)
    ax.plot(maxr * cos(r),maxr * sin(r),'k',lw=1.5,label=None)
    if zmin != 0:
        ax.plot(minr * cos(r),minr * sin(r),'k',lw=1.5,label=None)
    ax.axis('equal')

    # QSO and absorber positions
    for qso,theta,qname,l in zip(qsos,qtheta,'ABC',[la,lb,lc]):
        theta = (theta - theta0) * magnify_theta * RAD_PER_DEG
        qr = cosmo.Dc(min(qso.z,zmax)) * const
        
        dcol = dict(A='b',B='r',C='g')

        # custom vertex, rotate depending on theta
        vr = 1/np.sqrt(np.pi)
        va = theta+np.pi/2.
        vert = (vr*cos(va),vr*sin(va)),(-vr*cos(va),-vr*sin(va))
        print vert, vr, va
        # qso positions
        xq = qr * cos(theta)
        yq = qr * sin(theta)
        x2 = max([minr,wlosmin[1]]) * cos(theta)
        y2 = max([minr,wlosmin[1]]) * sin(theta)
        ax.plot([x2,xq],[y2,yq],'',lw=1.2,color=dcol[qname],label=None)
        if False:      #     qname != 'C':
            x0 = max([minr,wlosmin[0]]) * cos(theta)
            y0 = max([minr,wlosmin[0]]) * sin(theta)
            x1 = max([minr,wlosmax[0]]) * cos(theta)
            y1 = max([minr,wlosmax[0]]) * sin(theta)
            ax.plot([x0,x1],[y0,y1],'',lw=1.2,color=dcol[qname],label=None)

        label = 'QSO' if qname == 'A' else None
        ax.plot([xq],[yq],'+',ms=12,mew=2,color=dcol[qname],label=label)
        ax.text(xq-40,yq+30,qname,fontsize=18,zorder=100)

        # lya line positions
        l = l[(zmin < l.zabs)&(l.zabs < zmax)]
        l1 = l[l.ew>=1]
        l2 = l[l.ew<1]
        lr1 = np.array([cosmo.Dc(line.zabs) * const for line in l1])
        lr2 = np.array([cosmo.Dc(line.zabs) * const for line in l2])
        xl1 = lr1 * cos(theta)
        yl1 = lr1 * sin(theta)
        xl2 = lr2 * cos(theta)
        yl2 = lr2 * sin(theta)
        #label1 = ('Ly-$\\alpha$, ew $>= 1 \ \\AA$' if qname=='A' else None)
        #label2 = ('Ly-$\\alpha$, ew $< 1 \ \\AA$' if qname=='A' else None)
        #ax.plot(xl1,yl1,'^',mew=1,mec=dcol[qname],mfc=dcol[qname],label=label1)
        #ax.plot(xl2,yl2,'^',mew=1,mec=dcol[qname],mfc='None',label=label2)
        ax.scatter(xl1,yl1,marker=None,s=(60,),
                   edgecolor=dcol[qname],verts=vert,zorder=99)
        ax.scatter(xl2,yl2,marker=None,s=(15,),
                   edgecolor=dcol[qname],verts=vert,zorder=99)
        # selected metal line positions
        if metals:
            if qname=='A':
                c4 = np.array([cosmo.Dc(item)*const for item in
                               (0.2272,0.3995)])
                plot(c4 * cos(theta), c4 * sin(theta),'s',mew=1.5,
                     mfc='None',label='CIV')
            elif qname=='B':
                o6 = cosmo.Dc(0.2268) * const
                plot([o6*cos(theta)], [o6*sin(theta)],'o',mew=1.5,
                     mfc='None',label='OVI')
                lls = cosmo.Dc(0.3997)*const
                plot([lls*cos(theta)], [lls*sin(theta)],'^',mew=1.5,
                     mfc='None',label='LLS')
            elif qname=='C':
                lls = cosmo.Dc(0.5569) * const
                plot([lls*cos(theta)], [lls*sin(theta)],'^',mew=1.5,
                     mfc='None',label=None)

    # plot the position of each galaxy
    x = losdist * cos(sepdeg * RAD_PER_DEG)
    y = losdist * sin(sepdeg * RAD_PER_DEG)
    c0 = gal.isdeimos
    ax.plot(x[c0],y[c0],'ok', alpha=1, label='DEIMOS galaxy',zorder=1)
    ax.plot(x[~c0],y[~c0],'ok', mfc='None', alpha=1, label='CFHT galaxy',zorder=1)
    
    #ax.plot(x,y,'o',mfc=None,mec='b',mew=1.2,label='Galaxy')
    #ax.plot(x,y,'+b',mew=1.2,ms=10,label='Galaxy')
    
    if legend: ax.legend(loc='lower left',numpoints=1)


    xmin,xmax,ymin,ymax = ax.axis()
    # tick length
    ticklen = (xmax - xmin) * 0.008
    print 'ticklength', ticklen, 'ymin,ymax',ymin,ymax
    
    # r tickmarks and labels
    zrange = np.linspace(zmin,zmax,11)           # tuned for zmin=0,zmax=1
    tick_top = [cosmo.Dc(z) * const for z in zrange]
    tick_bottom = np.linspace(minr,maxr,11)
    tick_theta = maxsep * RAD_PER_DEG
    for rup,rdown,z in zip(tick_top[1:-1],tick_bottom[1:-1],zrange[1:-1]):
        foo = np.linspace(maxsep*RAD_PER_DEG,-maxsep*RAD_PER_DEG,100)
        ax.plot(rup * cos(foo),rup * sin(foo),'k:',lw=1,alpha=0.5,label=None)
        x1 = rup * cos(tick_theta)
        y1 = rup * sin(tick_theta)
        norm_theta = tick_theta+pi/2.
        xt = ticklen * cos(norm_theta)
        yt = ticklen * sin(norm_theta)
        ax.plot([x1,x1+xt],[y1,y1+yt],'k',lw=1.5,label=None)
        # top tick labels
        # find label offsets
        xt = 2.*ticklen * cos(norm_theta)
        yt = 2.*ticklen * sin(norm_theta)
        ax.text(x1+xt,y1+yt,'%3.1f' % z,rotation=maxsep,
                ha='center',fontsize=14)
        # bottom ticks
        x2 = rdown * cos(-tick_theta)
        y2 = rdown * sin(-tick_theta)
        xt = ticklen * cos(-norm_theta)
        yt = ticklen * sin(-norm_theta)
        ax.plot([x2,x2+xt],[y2,y2+yt],'k',lw=1.5,label=None)
        # bottom tick labels
        # find label offsets
        xt = 5.5*ticklen * cos(-norm_theta)
        yt = 5.5*ticklen * sin(-norm_theta)
        ax.text(x2+xt,y2+yt,'%4.0f' % rdown,rotation=-maxsep,
                ha='center',fontsize=14)

    # theta tickmarks and labels. One tick every 2 arcmin.
    tick_r = maxr
    tick_theta = np.arange(0,maxsep,2./60.*200)
    for i,t in enumerate(tick_theta):
        trad = t * RAD_PER_DEG
        x1 = tick_r * cos(trad)
        y1 = tick_r * sin(trad)
        xt = ticklen * cos(trad)
        yt = ticklen * sin(trad)
        ax.plot([x1,x1+xt],[y1,y1+yt],'k',lw=1.5,label=None)
        ax.plot([x1,x1+xt],[-y1,-(y1+yt)],'k',lw=1.5,label=None)
        temp = i*2./60.
        offset = 3*ticklen * cos(trad)        # tick label offset
        if deg:
            c1 = temp * RAD_PER_DEG
            c2 = -temp * RAD_PER_DEG
            ax.text(x1+offset,y1,'%.1e' % c1,va='center')
            if i != 0:  
                ax.text(x1+offset,-y1,'%.1e' % c2,va='center')            
        elif mode == 'ra' and not deg:
            c1 = dec2s(theta0 + temp,0)[0][:8].replace(' ',':')
            c2 = dec2s(theta0 - temp,0)[0][:8].replace(' ',':')
            ax.text(x1+offset,y1,c1,va='center')
            if i != 0:  
                ax.text(x1+offset,-y1,c2,va='center')
        elif mode == 'dec' and not deg:
            c1 = dec2s(0,theta0 + temp)[1][:9].replace(' ',':')
            c2 = dec2s(0,theta0 - temp)[1][:9].replace(' ',':')
            ax.text(x1+offset,y1,c1,va='center')
            if i != 0:  
                ax.text(x1+offset,-y1,c2,va='center')

    # label the axes
    xpos = (minr+0.5*(maxr-minr)) * cos(maxsep * RAD_PER_DEG)
    ypos = (minr+0.5*(maxr-minr)) * sin(maxsep * RAD_PER_DEG)
    offset = 10 * ticklen
    ax.text(xpos,ypos+offset-1.*ticklen,'Redshift',rotation=maxsep,fontsize=16,
            ha='center',va='center')
    ax.text(xpos,-ypos-offset,'LOS comoving distance (h$^{-1}$ Mpc)',
            rotation=-maxsep,fontsize=16,ha='center',va='center')
    if mode == 'ra':
        if deg:
            ax.text(maxr+ticklen*15,0,'RA sep (radians)',
                    rotation=-90,fontsize=16,ha='center',va='center')
        else:
            ax.text(maxr+ticklen*12,0,'Right Ascension (h:m:s)',
                    rotation=-90,fontsize=16,ha='center',va='center')
    elif mode == 'dec':
        if deg:
            ax.text(maxr+ticklen*15,0,'Dec sep (radians)',
                    rotation=-90,fontsize=16,ha='center',va='center')
        else:
            ax.text(maxr+ticklen*13,0,'Declination (d:m:s)',
                    rotation=-90,fontsize=16,ha='center',va='center')        

    # turn off the default tick marks and labels
   
    pl.setp(ax.get_xticklabels(),visible=False)
    pl.setp(ax.get_yticklabels(),visible=False)
    pl.setp(ax.get_xticklines(),visible=False)
    pl.setp(ax.get_yticklines(),visible=False)

    pl.show()
    
    return ax

if 0:
    
    gal = q0107.readgal()
    qsos = q0107.readqsos()
    la_all = q0107.readabs('A')
    lb_all = q0107.readabs('B')
    lc_all = q0107.readabs('C')
    maxwavs = wlya*(1 + qsos.z)*(1. - 3500./Ckms) # 3500 km/s from lya emission
    # get lya only
    lyalines = []
    for i,l in enumerate([la_all,lb_all,lc_all]):
        lya_ids = ('unknown','lya')
        cond = np.array([ID in lya_ids for ID in l.ID])
        cond = cond & (l.wav < maxwavs[i])
        cond = cond & (l.signif_det > NSIGMA)
        temp = l[cond]
        temp.zabs = temp.wav/wlya - 1.
        lyalines.append(temp)

    la,lb,lc = lyalines

    #from gas_gal_hiz import read
    #lbgs = read.lbgs('he0940')
    zmin,zmax=0.0,1.1
    #ra,dec = s2dec('09 42 53.5','-11 04 25.9')
    #qsos = (np.rec.fromrecords([(ra,dec,3.053)],names='ra,dec,z'))
    #pie(lbgs,qsos,mode='dec',deg=0,zmin=zmin,zmax=zmax,magnify_theta=10,lines=None)


    # deimos gals
    z,ra,dec,mag = readcol('/home/nhmc/proposals/keck/09a/all.redshifts',
                           com='#',col=(1,2,3,5))
    ra,dec = zip(*[s2dec(r,d) for r,d in zip(ra,dec)])
    deimos = np.rec.fromarrays([ra,dec,z,mag,np.ones(len(z),bool)],
                               names='ra,dec,z,mag,isdeimos')
    temp = [gal[k] for k in 'ra dec z rmag'.split()]
    temp = np.rec.fromarrays(temp + [np.zeros(len(gal),bool)],
                             names='ra,dec,z,mag,isdeimos')
    allgal = np.concatenate([temp,deimos]).view(np.recarray)

    pie(allgal,qsos,mode='dec',deg=0,zmin=zmin,zmax=zmax,
        magnify_theta=200,legend=0)
    pie(allgal,qsos,mode='ra',deg=0,zmin=zmin,zmax=zmax,
        magnify_theta=200)

